import json
import os
from client import MCPClient
import asyncio
import websockets
client = MCPClient()
async def handle_client(websocket, path):
    try:
        cookie = path.split('=')[1]
        os.environ['FUEL_TOKEN'] = cookie
        servers = {
            "server": "server.py"
        }
        await client.connect_to_server(servers)
        message = await websocket.recv()
        print(f"Received from client on /query: {message}")
        
        response_result = await client.api_call(message)
        print("🚀 ~ response_result:", response_result, response_result['tool_name'])
        if response_result['tool_name'] == 'serve_jump_page':
            content = {
                "val": response_result['result'],
                "tool_name": response_result['tool_name'],
                "isDone": True
            }
            await websocket.send(json.dumps(content, ensure_ascii=False))
        else:
            for chunk in response_result['result']:
                if hasattr(chunk, 'choices'):
                    choice = chunk.choices[0]
                    if hasattr(choice, 'delta') and hasattr(choice.delta, 'content'):
                        print("🚀 ~ choice:", choice)
                        # 定义一个对象
                        content = {
                            "val": choice.delta.content,
                            "tool_name": response_result['tool_name'],
                            "isDone": choice.finish_reason == "stop"
                        }
                        if "chat_url" in response_result:
                            content["chat_url"] = response_result['chat_url']
                        await websocket.send(json.dumps(content, ensure_ascii=False))
    except websockets.exceptions.ConnectionClosedOK:
        pass
async def main():
    async with websockets.serve(handle_client, "localhost", 8765) as server:
        await server.wait_closed()

if __name__ == '__main__':
    print("🚀 ~ localhost://8765/query:")
    asyncio.run(main())